"""

Visualization tools.

"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d


def vis_vectors_2d(data, label):
    """
    Visualize the first and second dimension of vectors in the data.
    
    Args:
    - data: :list: (n,) list of vectors.
    - label: :list: (n,) list of corresponding labels.
    """
    types = ['gufeng', 'liuxing', 'minyao', 'shuochang']
    color = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']
    
    plt.figure()
    for i in range(len(data)):
        plt.scatter(data[i][0], data[i][1], c=color[label[i]])
    plt.show()
    
def vis_vectors_3d(data, label):
    """
    Visualize the first and second dimension of vectors in the data.
    
    Args:
    - data: :list: (n,) list of vectors.
    - label: :list: (n,) list of corresponding labels.
    """
    types = ['gufeng', 'liuxing', 'minyao', 'shuochang']
    color = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']

    plt.figure("3D Scatter", facecolor="lightgray")
    ax3d = plt.gca(projection="3d")
    ax3d.set_xlabel('x', fontsize=14)
    ax3d.set_ylabel('y', fontsize=14)
    ax3d.set_zlabel('z', fontsize=14)

    for i in range(len(data)):
        ax3d.scatter(data[i][0], data[i][1], data[i][2], c=color[label[i]])
    plt.show()
    
def vis_loss(train_losses, val_losses):
    """
   Visualize the loss during training. 
    """
    iter_per_ep = len(train_losses) / (len(val_losses) - 1) 
    train_indices = list(range(len(train_losses)))
    val_indices = np.array(range(len(val_losses))) * iter_per_ep
    
    plt.figure()
    plt.plot(train_indices, train_losses)
    plt.plot(val_indices, val_losses)
    plt.legend(['train_loss','val_loss'])
    plt.show()
       
    
    
    
    
